package com.rizzo.back.helper;

import com.aliasi.lm.TokenizedLM;
import com.aliasi.tokenizer.TokenizerFactory;
import com.aliasi.util.ScoredObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.lang.reflect.InvocationTargetException;
import java.util.ArrayList;
import java.util.List;
import java.util.SortedSet;

/**
 * TODO
 */
public class SipHelper {

    private static final Logger LOGGER = LoggerFactory.getLogger(SipHelper.class);

    private TokenizerFactory tokenizerFactory;

    public List<String> findSIPs(List<String> backgroundData, List<String> foregroundData, int tokenizerNgram,
                                 int pruneFactor, int reportNgram, int minCount, int maxReturned)
            throws IOException, InvocationTargetException, NoSuchMethodException, IllegalAccessException {
        List<String> sips = new ArrayList<String>();
        // training background model
        TokenizedLM backgroundModel = trainModel(tokenizerFactory, tokenizerNgram, backgroundData);
        backgroundModel.sequenceCounter().prune(pruneFactor);
        // training foreground model
        TokenizedLM foregroundModel = trainModel(tokenizerFactory, tokenizerNgram, foregroundData);
        foregroundModel.sequenceCounter().prune(pruneFactor);
        // assembling new terms in test vs. training
        SortedSet<ScoredObject<String[]>> newTerms = foregroundModel.newTermSet(reportNgram, minCount, maxReturned, backgroundModel);
        for (ScoredObject<String[]> newTerm : newTerms) {
            double score = newTerm.score();
            String[] tokens = newTerm.getObject();
            StringBuffer phrase = new StringBuffer();
            for (String token : tokens) {
                phrase.append(token).append(" ");
            }
            LOGGER.debug("Score: " + score + " - Phrase: " + phrase);
            sips.add(phrase.toString());
        }
        return sips;
    }

    private TokenizedLM trainModel(TokenizerFactory tokenizerFactory, int tokenizerNgram, List<String> backgroundData) {
        TokenizedLM model = new TokenizedLM(tokenizerFactory, tokenizerNgram);
        for (String dataEntry : backgroundData) {
            model.train(dataEntry);
        }
        return model;
    }

    public void setTokenizerFactory(TokenizerFactory tokenizerFactory) {
        this.tokenizerFactory = tokenizerFactory;
    }
}
